refactor flush data worker - #470

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@build_tx_service.cmake`:
- Around line 33-43: The ASAN detection branch using CMAKE_CXX_FLAGS and the
conditional find_library(NAMES boost_context-asan) is dead and should be removed
or replaced with an explicit ASAN option; implement ASAN support by adding a
boolean option (e.g., USE_ASAN), then apply
add_compile_options(-fsanitize=address) / add_compile_definitions(...) as done
in build_eloq_store.cmake and conditionally pick the ASAN library when USE_ASAN
is ON; also reorder the logic so find_package(Boost 1.70 REQUIRED) runs before
calling find_library(Boost_CONTEXT_LIBRARY) and pass Boost_LIBRARY_DIRS (or
Boost::boost targets) as hints to find_library, and avoid a single plain
${Boost_CONTEXT_LIBRARY} for multi-config generators (use generator expressions
if conditional linking is needed later).

In `@store_handler/data_store_service_client_closure.cpp`:
- Around line 525-583: The loops in SyncPutAllData::Wait,
SyncConcurrentRequest::WaitForCapacityAndIncrement, and
SyncConcurrentRequest::WaitForAll set waiting_ then unlock and call yield_fn(),
which risks a lost-wakeup if a completion fires resume_fn before the coroutine
actually suspends; fix by adding a handshake atomic (e.g., suspended_) and
change the wait loops to set suspended_.store(true, ...) immediately before
unlocking and calling (*yield_fn_)(), and change the completion/resume logic to
only call resume_fn when it successfully observes-and-clears suspended_ (use
compare_exchange to set suspended_ to false when delivering the wakeup); ensure
waiting_ is still used for diagnostics but the actual resume decision is gated
by suspended_ so resume cannot be lost.

In `@tx_service/include/cc/cc_req_misc.h`:
- Around line 897-914: The coroutine handoff can miss a resume because waiting_
is armed before the coroutine actually yields; fix by introducing an atomic
"resume_pending_" (or wakeup_pending_) flag and update the handshake: have
AbortCcRequest/Execute/SetFinished set resume_pending_.store(true) when they
would call resume_fn_, and in WaitableCc::Wait and UpdateCceCkptTsCc::Wait
check/clear resume_pending_ under lock before attempting to yield and only arm
waiting_ while holding the lock so the resume path cannot race; specifically,
change the loop that currently does waiting_.store(true); lk.unlock();
(*yield_fn)(); lk.lock(); waiting_.store(false); to: under lk check if
resume_pending_ is true (consume it and continue), otherwise set waiting_=true
while holding lk, then unlock and call (*yield_fn)(); on re-lock clear waiting_,
and ensure resume code sets resume_pending_ and invokes resume_fn_ if waiting_
is observed or just sets resume_pending_ otherwise so the waiter will see the
pending resume before yielding.

In `@tx_service/include/cc/cc_request.h`:
- Around line 8979-8987: The multiplication in UpdatedMemory() can overflow
because dirty_key_count_ * memory_usage_ is done in the narrower type; fix by
performing the multiplication in a wider integer type (e.g., uint64_t or
unsigned long long) before dividing and then cast the final result back to
size_t. Update the expression in UpdatedMemory() to cast dirty_key_count_ (or
memory_usage_) to the wider type for the intermediate, perform the arithmetic
and rounding there, and return static_cast<size_t>(...) to avoid silent
underestimation while keeping the function signature and logic the same.

In `@tx_service/include/store/data_store_handler.h`:
- Around line 87-95: Derived overrides (PutAll, PutArchivesAll,
CopyBaseToArchive) in the Bigtable/Dynamo handlers currently use different
parameter types (e.g., vector<FlushRecord>, vector<pair<TxKey,int32_t>>) and
thus don't match the base virtual signatures; update these methods in the
derived classes to exactly match the base declarations:
PutAll(...unordered_map<string_view,
vector<unique_ptr<txservice::FlushTaskEntry>>> &flush_task, const
function<void()> *yield_fptr = nullptr, const function<void()> *resume_fptr =
nullptr, const function<void()> *sync_yield_fptr = nullptr), and make
PutArchivesAll and CopyBaseToArchive take the same parameter types as the base
(replace any FlushRecord/vector<pair<...>> parameters with the unordered_map of
unique_ptr<txservice::FlushTaskEntry>), remove or change any conflicting types
like FlushRecord to use FlushTaskEntry, and add the missing UpdateRangeSlices
overload that accepts vector<UpdateRangeSlicesReq> plus the callback parameters
so the derived classes implement both base overloads.

In `@tx_service/src/cc/local_cc_shards.cpp`:
- Around line 6141-6213: The loop currently prioritizes pending_flush_work over
resume_queue causing resumed coroutines (CoroCtx) to starve; change the dispatch
order in the flush worker loop to check and resume entries from resume_queue
(resume_queue.front(), ctx->coro_.resume(), etc.) before pulling new work from
pending_flush_work, or otherwise drain/round-robin resume_queue at least once
per loop iteration so resumed coroutines get scheduled prior to starting fresh
FlushDataImpl tasks; update references in this block that use
pending_flush_work, resume_queue, ctx->coro_, FlushDataImpl, and CoroCtx
accordingly.
- Around line 4811-4819: partition_ids is computed round-robin
(partition_number_this_core and partition_ids via worker_idx + core_number * i)
but later scans use a contiguous [min,max] filter causing workers to overlap
when partition_number_per_scan > 1; change the scan filter to match the actual
round-robin set instead of a single contiguous range: for the scan code that
consumes partition_ids (the logic that builds the [min,max] filter), replace it
with either an explicit per-partition filter (iterate partition_ids and issue
sub-scans or add multiple non-contiguous ranges) or build filters with stride
core_number (e.g., create ranges or a predicate that accepts id % core_number ==
worker_idx) so that scans only cover the partitions in partition_ids and no
duplicates are exported (ensure this change is applied to the similar block
around partition_number_this_core usage at the other location as well).
- Around line 5161-5213: The code filters out records with rec.cce_ == nullptr
when building data_sync_vec (using scan_cc.DataSyncVec()), but
scan_cc.ArchiveVec() and scan_cc.MoveBaseIdxVec() contain indices into the
original unfiltered DataSyncVec, so those indices become invalid; to fix this
either (A) perform the filtering (build data_sync_vec and mv_base_vec) before
ExportForCkpt generates ArchiveVec()/MoveBaseIdxVec() so the export sees
compacted indices, or (B) if export cannot change, create a stable index-mapping
when you build data_sync_vec (map original DataSyncVec() index -> new
data_sync_vec index) and then update the entries in scan_cc.ArchiveVec() and
scan_cc.MoveBaseIdxVec() to use the remapped indices (or remap keys via that map
when calling rec.SetKey and when emplacing mv_base_vec), ensuring all references
(scan_cc.ArchiveVec(), scan_cc.MoveBaseIdxVec(), rec.GetKeyIndex(),
rec.SetKey(), and places using TxKey key_raw) point to the compacted positions.

In `@tx_service/src/tx_index_operation.cpp`:
- Around line 465-488: The log unconditionally dereferences cluster_config_addr
after re-scanning MetaDataReadSet(), which can be nullptr and crash; update the
unlock failure logging in the block that references cluster_config_addr (and the
loop that sets it) to first check cluster_config_addr != nullptr and only append
Term(), NodeGroupId(), and Sharder::Instance().LeaderTerm(...) when non-null,
otherwise log a safe placeholder (or reuse the validated address stored when
unlock_cluster_config_op_.Reset(...) was built); ensure references to
cluster_config_addr, meta_rset, and unlock_cluster_config_op_.hd_result_ are
guarded to prevent null dereference.

---

Outside diff comments:
In `@store_handler/data_store_service_client_closure.h`:
- Around line 101-151: The coroutine-enabled subclasses (FetchTableCallbackData,
FetchDatabaseCallbackData, FetchAllDatabaseCallbackData,
UpsertDatabaseCallbackData, DropDatabaseCallbackData,
DiscoverAllTableNamesCallbackData) still implement one-shot overrides of
Wait()/Notify() which lose a wakeup if Notify() happens before the waiter
yields; update each subclass to use the same finished_ + waiting_ handshake used
by the base SyncCallbackData: remove or refactor their custom Wait()/Notify()
overrides to call the base implementations (or replicate the
waiting_.store/finished_ checks and memory_order semantics used in the base
SetCoroCallbacks/Wait/Notify), ensure SetCoroCallbacks is used to store the
yield/resume pointers, and preserve unlocking before invoking resume_fn_ so the
resume path follows the base-class pattern to avoid lost wakeups.

In `@store_handler/data_store_service_client.cpp`:
- Around line 3120-3150: The final BatchWriteRecords call dispatches before
incrementing sync_concurrent->unfinished_request_cnt_, risking a race where the
callback decrements a not-yet-incremented counter; replace the manual increment
block before that dispatch with the same capacity-guarded increment used
elsewhere by calling WaitForCapacityAndIncrement(sync_concurrent) (or the
existing wrapper used in this file) so the in-flight counter is incremented
before BatchWriteRecords(kv_mvcc_archive_name, ...) is invoked; remove the
explicit unique_lock and unfinished_request_cnt_++ block and ensure
SyncConcurrentRequestCallback and WaitForAll() semantics remain unchanged.

In `@tx_service/include/cc/cc_req_misc.h`:
- Around line 871-879: The Reset implementation for WaitableCc (the one calling
RunOnTxProcessorCc::Reset and setting unfinished_cnt_ and error_code_) must also
clear the previous callback/wait state to avoid carrying stale coroutine wakeup
paths: explicitly reset yield_fn_ and resume_fn_ to empty/null and set waiting_
to false so any recycled request cannot resume a prior Wait; apply the same
clearing changes in the other Reset overload (the similar block around the other
Reset that mirrors these fields).

In `@tx_service/include/cc/local_cc_shards.h`:
- Around line 474-494: The WITH_JEMALLOC initialization runs unconditionally and
can overwrite hash_partition_ckpt_arena_id_ on subsequent calls; move the entire
jemalloc branch (mallctl creating the arena and setting
hash_partition_ckpt_arena_id_) inside the guarded block that checks/creates
hash_partition_ckpt_heap_ within InitializeHashPartitionCkptHeap(), so the
jemalloc arena is created only once alongside mi_heap_new and
hash_partition_main_thread_id_, referencing symbols
InitializeHashPartitionCkptHeap, hash_partition_ckpt_heap_,
hash_partition_ckpt_arena_id_, hash_partition_main_thread_id_, mi_heap_new and
mallctl to locate the correct code to relocate.
- Around line 2015-2023: The callbacks yield_fptr/resume_fptr/sync_yield_fptr
passed into PostProcessFlushTaskEntries (and forwarded into
UpdateRangeSlices/PutAllImpl) are currently stored as raw pointers inside pooled
types SyncPutAllData, SyncConcurrentRequest and SyncCallbackData, creating a
lifetime/USE-AFTER-FREE risk; modify those pooled objects and all call sites so
they take ownership of the callbacks (either by copying the std::function into a
member std::function or by storing a std::shared_ptr<std::function<...>>),
update constructors/factory methods that create
SyncPutAllData/SyncConcurrentRequest/SyncCallbackData to accept the owning form,
and ensure callers to PostProcessFlushTaskEntries/UpdateRangeSlices/PutAllImpl
pass owned callbacks (not pointers) so completion callbacks can safely invoke
them after async/pooled lifetimes.

---

Nitpick comments:
In `@store_handler/data_store_service_client.cpp`:
- Around line 318-320: Add a boundary check that fails fast when coroutine
callback pointers are half-configured: ensure yield_fptr and resume_fptr are
either both null or both non-null (and similarly handle any sync pair like
sync_yield_fptr with its companion where applicable). Implement this as a small
helper (e.g., ValidateCoroutineCallbacks or a DCHECK pair check) and call it at
the start of the entry points that accept these pointers (the functions taking
yield_fptr, resume_fptr, sync_yield_fptr) so miswired integrations immediately
trigger an assertion/error; apply the same check to the other listed call sites
(lines ~558-559, ~1792-1794, ~2950-2951, ~3191-3192).

In `@store_handler/rocksdb_handler.cpp`:
- Around line 526-653: The coroutine branch duplicates the full RocksDB write
loop (inside the SubmitWork lambda) causing TTL/delete handling, metrics, and
error paths to diverge from the synchronous PutAll/PersistKV code; refactor by
extracting the core write logic into a shared helper (e.g., a new method like
ExecutePutAllBatch or PersistBatchToRocksDB) that accepts the DB pointer
(GetDBPtr()), write options, the batch structure, batch_write_size_, and a
callback/result indicator, and then have both the coroutine path (the lambda
submitted via query_worker_pool_->SubmitWork and using
RocksDBWriteSyncCallback/Wait) and the original synchronous path call that
helper so all TTL checks (EncodeToKvKey, SerializeFlushRecord, RecordStatus/TTL
logic), metrics::kv_meter collection, error logging
(status.ToString()/status.code()), and write_batch management live in one place;
ensure the helper returns success/failure so the coroutine lambda invokes
callback.Notify(...) accordingly and the synchronous callers propagate the same
result.

In `@tx_service/include/cc/local_cc_shards.h`:
- Around line 2546-2551: The AddFlushTaskEntry() call currently infers the
target flush worker from entry->data_sync_task_->id_, which hides the mapping
when data_sync and flush worker counts differ; change AddFlushTaskEntry()
signature to accept an explicit worker identifier (either the source
data_sync_worker_id or the resolved flush_worker_id) and update all call sites
to pass DataSyncWorkerToFlushDataWorker(data_sync_worker_id) or the original
data_sync_worker_id as appropriate; update the implementation of
AddFlushTaskEntry(), its callers, and any tests to use the new explicit
parameter so buffer/queue selection is clear and testable, and remove implicit
reliance on entry->data_sync_task_->id_ inside AddFlushTaskEntry().

In `@tx_service/include/cc/template_cc_map.h`:
- Around line 8125-8147: AdjustDataKeyStats currently updates data_key_count_
and dirty_data_key_count_ but doesn't validate the invariant
dirty_data_key_count_ <= data_key_count_; add a cheap assertion at the end of
AdjustDataKeyStats (after the existing updates and after the early return for
table_name_.IsMeta()) that checks dirty_data_key_count_ is never greater than
data_key_count_ (taking care of the size_t vs signed conversions) so future
missed bookkeeping fails fast; reference AdjustDataKeyStats, data_key_count_,
dirty_data_key_count_, and table_name_.IsMeta() to locate where to insert the
assertion.

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@tx_service/include/cc/cc_req_misc.h`:
- Around line 897-914: The Wait(...) coroutine handoff races because waiting_ is
set before the coroutine actually yields, allowing a pre-yield resume (from
AbortCcRequest, Execute, SetFinished or where resume_fn_/resume logic runs) to
be lost; add a protected "pending_resume" boolean/state guarded by mux_ that is
set by the producers instead of directly invoking resume_fn_ when the waiter
hasn't yet yielded, and modify Wait (and the other Wait variants at the same
pattern) to check and clear pending_resume while holding mux_ before actually
unlocking and calling (*yield_fn)(), and only call (*yield_fn)() when no pending
resume remains (if pending_resume is set, consume it and call (*resume_fn)()
instead of yielding). Ensure producers set pending_resume under mutex (or invoke
resume_fn_ immediately if waiting_ is true and pending_resume is clear), and
clear pending_resume when consumed so no resume is lost.

In `@tx_service/src/cc/local_cc_shards.cpp`:
- Around line 6141-6213: The loop currently favors starting new FlushData tasks
(pending_flush_work) before resuming coroutines (resume_queue), causing resumed
coroutines to starve; reorder the branch checks so that the resume_queue
handling runs before pending_flush_work, i.e., check and pop from resume_queue
and call ctx->coro_.resume() (using the existing resume_queue,
ctx->coro_.resume(), and flush_worker_lk logic) prior to creating/launching a
new coro from pending_flush_work, and ensure the same lock/unlock semantics
around flush_worker_lk are preserved when swapping the two blocks so lock
handling and notify_all calls remain correct.
- Around line 4856-4875: The filter_lambda incorrectly treats the round-robin
selection as a contiguous range using min_partition_id_this_scan and
max_partition_id_this_scan; instead, change the predicate in the lambda
(filter_lambda) to test membership against the actual partition_ids slice for
this scan (from index i to i + partition_number_per_scan - 1) rather than using
a range check — e.g., capture the relevant partition_ids sub-range (or build a
small local set/vector of those IDs) and check (hash_code %
total_hash_partitions) is equal to one of those IDs in addition to honoring
data_sync_task->filter_lambda_; update the captures accordingly so partition_ids
and partition subset are available inside filter_lambda.
- Around line 5161-5213: The archive/move-base loops use original scan positions
but earlier code compacts out records with null cce_, so create a remapping from
original scan indices to compaction indices: allocate a vector<int> (size
scan_cc.accumulated_scan_cnt_) initialized to -1, and in the loop that emplaces
into data_sync_vec (the block using scan_cc.DataSyncVec(), rec.cce_, and
data_sync_vec->emplace_back) set map[j] = current_compacted_index (increment as
you push); then in the ArchiveVec loop
(rec.SetKey((*data_sync_vec)[rec.GetKeyIndex()].Key())) and MoveBaseIdxVec
handling (key_idx = scan_cc.MoveBaseIdxVec()[j]; TxKey key_raw =
(*data_sync_vec)[key_idx].Key()) replace uses of the original index with the
mapped index (map[orig_idx]) and handle map==-1 (skip or assert/error) to avoid
wrong keys or OOB access. Ensure you reference scan_cc.DataSyncVec(),
scan_cc.ArchiveVec(), scan_cc.MoveBaseIdxVec(), data_sync_vec->emplace_back, and
FlushRecord::SetKey when applying the remap.

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tx_service/include/cc/local_cc_shards.h`:
- Around line 476-481: The arenas.create call and assignment to
hash_partition_ckpt_arena_id_ must be moved inside the one-time init guard so it
only runs when hash_partition_ckpt_heap_ is null; currently arenas.create runs
on every invocation and overwrites hash_partition_ckpt_arena_id_, leaking prior
arenas. Update the init sequence around hash_partition_ckpt_heap_ and
hash_partition_main_thread_id_ (the if (!hash_partition_ckpt_heap_) block) to
also call arenas.create and set hash_partition_ckpt_arena_id_ there, ensuring
subsequent calls skip arena creation and do not mutate or leak the previous
arena ID.

In `@tx_service/include/fault/log_replay_service.h`:
- Around line 230-232: The map currently stores std::unique_ptr<ConnectionInfo>
in inbound_connections_ which doesn't prevent use-after-free when callbacks
(on_received_messages, on_idle_timeout, on_closed) grab raw pointers under
inbound_mux_ and use them after unlocking; change the ownership to
std::shared_ptr<ConnectionInfo> (i.e., make inbound_connections_ an
unordered_map<brpc::StreamId, std::shared_ptr<ConnectionInfo>>) and update all
places that construct entries to create shared_ptrs; then in each callback
acquire a local std::shared_ptr<ConnectionInfo> under inbound_mux_ (copying the
map value) and use that local shared_ptr after releasing the lock to keep the
object alive.

---

Duplicate comments:
In `@tx_service/include/cc/cc_req_misc.h`:
- Around line 881-886: Set up the full coroutine wait/resume handshake by
ensuring Wait installs the callbacks into the members read by
AbortCcRequest()/Execute(): call SetCoroCallbacks(...) (or assign
yield_fn_/resume_fn_ directly) before any chance to be woken, only arm waiting_
after the coroutine actually yields (i.e., set waiting_ from inside the yield
function or immediately after calling the yield callback), and clear waiting_,
yield_fn_ and resume_fn_ when resetting/reusing a WaitableCc (e.g., in
Reset()/constructor) so no stale pointers remain; apply the same changes to the
other Wait implementations referenced (lines ~897-914, 934-975, 998-1001) to
keep behavior consistent.
- Around line 1048-1053: The lost-wakeup race comes from SetCoroCallbacks/Wait
using separate local args and member resume_fn_/yield_fn_ and setting waiting_
before the published callbacks, so SetFinished can race and miss the resume; fix
by making Wait publish the coroutine callbacks into the shared members (assign
yield_fn_ and resume_fn_) before flipping waiting_, or alternatively have
SetFinished atomically snapshot and invoke any resume_fn_ even if waiting_
races; update functions UpdateCceCkptTsCc/Wait/SetCoroCallbacks/SetFinished to
either (a) assign yield_fn_ and resume_fn_ into the members before setting
waiting_ and use appropriate memory ordering, or (b) in SetFinished always load
resume_fn_ into a local under the same synchronization and call it if non-null,
ensuring no lost-wakeup.